// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "zencore.h"

#include <zencore/intmath.h>
#include <zencore/thread.h>

#include <cstddef>
#include <cstring>
#include <span>
#include <vector>

namespace zen {

#if defined(__cpp_lib_ranges) && __cpp_lib_ranges >= 201911L
template<typename T>
concept ContiguousRange = std::ranges::contiguous_range<T>;
#else
template<typename T>
concept ContiguousRange = true;
#endif

struct MemoryView;

class MemoryArena
{
public:
	ZENCORE_API MemoryArena();
	ZENCORE_API ~MemoryArena();

	ZENCORE_API void* Alloc(size_t Size, size_t Alignment);
	ZENCORE_API void  Free(void* Ptr);

private:
};

class Memory
{
public:
	ZENCORE_API static void* Alloc(size_t Size, size_t Alignment = sizeof(void*));
	ZENCORE_API static void	 Free(void* Ptr);
};

/** Allocator which claims fixed-size blocks from the underlying allocator.

	There is no way to free individual memory blocks.

	\note This is not thread-safe, you will need to provide synchronization yourself
*/

class ChunkingLinearAllocator
{
public:
	explicit ChunkingLinearAllocator(uint64_t ChunkSize, uint64_t ChunkAlignment = sizeof(std::max_align_t));
	~ChunkingLinearAllocator();

	ZENCORE_API void Reset();

	ZENCORE_API void* Alloc(size_t Size, size_t Alignment = sizeof(void*));
	inline void		  Free(void* Ptr) { ZEN_UNUSED(Ptr); /* no-op */ }

	ChunkingLinearAllocator(const ChunkingLinearAllocator&) = delete;
	ChunkingLinearAllocator& operator=(const ChunkingLinearAllocator&) = delete;

private:
	uint8_t*		   m_ChunkCursor	  = nullptr;
	uint64_t		   m_ChunkBytesRemain = 0;
	const uint64_t	   m_ChunkSize		  = 0;
	const uint64_t	   m_ChunkAlignment	  = 0;
	std::vector<void*> m_ChunkList;
};

//////////////////////////////////////////////////////////////////////////

struct MutableMemoryView
{
	MutableMemoryView() = default;

	MutableMemoryView(void* DataPtr, size_t DataSize)
	: m_Data(reinterpret_cast<uint8_t*>(DataPtr))
	, m_DataEnd(reinterpret_cast<uint8_t*>(DataPtr) + DataSize)
	{
	}

	MutableMemoryView(void* DataPtr, void* DataEndPtr)
	: m_Data(reinterpret_cast<uint8_t*>(DataPtr))
	, m_DataEnd(reinterpret_cast<uint8_t*>(DataEndPtr))
	{
	}

	inline bool IsEmpty() const { return m_Data == m_DataEnd; }
	void*		GetData() const { return m_Data; }
	void*		GetDataEnd() const { return m_DataEnd; }
	size_t		GetSize() const { return reinterpret_cast<uint8_t*>(m_DataEnd) - reinterpret_cast<uint8_t*>(m_Data); }

	inline bool EqualBytes(const MutableMemoryView& InView) const
	{
		const size_t Size = GetSize();

		return Size == InView.GetSize() && (memcmp(m_Data, InView.m_Data, Size) == 0);
	}

	/** Modifies the view to be the given number of bytes from the right. */
	inline void RightInline(uint64_t InSize)
	{
		const uint64_t OldSize = GetSize();
		const uint64_t NewSize = zen::Min(OldSize, InSize);
		m_Data				   = GetDataAtOffsetNoCheck(OldSize - NewSize);
		m_DataEnd			   = m_Data + NewSize;
	}

	/** Returns the right-most part of the view by taking the given number of bytes from the right. */
	[[nodiscard]] inline MutableMemoryView Right(uint64_t InSize) const
	{
		MutableMemoryView View(*this);
		View.RightChopInline(InSize);
		return View;
	}

	/** Modifies the view by chopping the given number of bytes from the left. */
	inline void RightChopInline(uint64_t InSize)
	{
		const uint64_t Offset = zen::Min(GetSize(), InSize);
		m_Data				  = GetDataAtOffsetNoCheck(Offset);
	}

	/** Returns the left-most part of the view by taking the given number of bytes from the left. */
	constexpr inline MutableMemoryView Left(uint64_t InSize) const
	{
		MutableMemoryView View(*this);
		View.LeftInline(InSize);
		return View;
	}

	/** Modifies the view to be the given number of bytes from the left. */
	constexpr inline void LeftInline(uint64_t InSize) { m_DataEnd = zen::Min(m_DataEnd, m_Data + InSize); }

	/** Modifies the view to be the middle part by taking up to the given number of bytes from the given offset. */
	inline void MidInline(uint64_t InOffset, uint64_t InSize = ~uint64_t(0))
	{
		RightChopInline(InOffset);
		if (InSize != ~uint64_t(0))
		{
			LeftInline(InSize);
		}
	}

	/** Returns the middle part of the view by taking up to the given number of bytes from the given position. */
	[[nodiscard]] inline MutableMemoryView Mid(uint64_t InOffset, uint64_t InSize = ~uint64_t(0)) const
	{
		MutableMemoryView View(*this);
		View.MidInline(InOffset, InSize);
		return View;
	}

	/** Returns the right-most part of the view by chopping the given number of bytes from the left. */
	[[nodiscard]] inline MutableMemoryView RightChop(uint64_t InSize) const
	{
		MutableMemoryView View(*this);
		View.RightChopInline(InSize);
		return View;
	}

	inline MutableMemoryView& operator+=(size_t InSize)
	{
		RightChopInline(InSize);
		return *this;
	}

	/** Copies bytes from the input view into this view, and returns the remainder of this view. */
	inline MutableMemoryView CopyFrom(MemoryView InView) const;

private:
	uint8_t* m_Data	   = nullptr;
	uint8_t* m_DataEnd = nullptr;

	/** Returns the data pointer advanced by an offset in bytes. */
	inline constexpr uint8_t* GetDataAtOffsetNoCheck(uint64_t InOffset) const { return m_Data + InOffset; }
};

//////////////////////////////////////////////////////////////////////////

struct MemoryView
{
	MemoryView() = default;

	MemoryView(const MutableMemoryView& MutableView)
	: m_Data(reinterpret_cast<const uint8_t*>(MutableView.GetData()))
	, m_DataEnd(m_Data + MutableView.GetSize())
	{
	}

	MemoryView(const void* DataPtr, size_t DataSize)
	: m_Data(reinterpret_cast<const uint8_t*>(DataPtr))
	, m_DataEnd(reinterpret_cast<const uint8_t*>(DataPtr) + DataSize)
	{
	}

	MemoryView(const void* DataPtr, const void* DataEndPtr)
	: m_Data(reinterpret_cast<const uint8_t*>(DataPtr))
	, m_DataEnd(reinterpret_cast<const uint8_t*>(DataEndPtr))
	{
	}

	inline bool Contains(const MemoryView& Other) const { return (m_Data <= Other.m_Data) && (m_DataEnd >= Other.m_DataEnd); }
	inline bool IsEmpty() const { return m_Data == m_DataEnd; }
	const void* GetData() const { return m_Data; }
	const void* GetDataEnd() const { return m_DataEnd; }
	size_t		GetSize() const { return reinterpret_cast<const uint8_t*>(m_DataEnd) - reinterpret_cast<const uint8_t*>(m_Data); }
	inline bool operator==(const MemoryView& Rhs) const { return m_Data == Rhs.m_Data && m_DataEnd == Rhs.m_DataEnd; }

	inline bool EqualBytes(const MemoryView& InView) const
	{
		const size_t Size = GetSize();

		return Size == InView.GetSize() && (memcmp(m_Data, InView.GetData(), Size) == 0);
	}

	inline MemoryView& operator+=(size_t InSize)
	{
		RightChopInline(InSize);
		return *this;
	}

	/** Modifies the view by chopping the given number of bytes from the left. */
	inline void RightChopInline(uint64_t InSize)
	{
		const uint64_t Offset = zen::Min(GetSize(), InSize);
		m_Data				  = GetDataAtOffsetNoCheck(Offset);
	}

	inline MemoryView RightChop(uint64_t InSize) const
	{
		MemoryView View(*this);
		View.RightChopInline(InSize);
		return View;
	}

	/** Returns the right-most part of the view by taking the given number of bytes from the right. */
	[[nodiscard]] inline MemoryView Right(uint64_t InSize) const
	{
		MemoryView View(*this);
		View.RightInline(InSize);
		return View;
	}

	/** Modifies the view to be the given number of bytes from the right. */
	inline void RightInline(uint64_t InSize)
	{
		const uint64_t OldSize = GetSize();
		const uint64_t NewSize = zen::Min(OldSize, InSize);
		m_Data				   = GetDataAtOffsetNoCheck(OldSize - NewSize);
		m_DataEnd			   = m_Data + NewSize;
	}

	/** Returns the left-most part of the view by taking the given number of bytes from the left. */
	inline MemoryView Left(uint64_t InSize) const
	{
		MemoryView View(*this);
		View.LeftInline(InSize);
		return View;
	}

	/** Modifies the view to be the given number of bytes from the left. */
	inline void LeftInline(uint64_t InSize)
	{
		InSize	  = zen::Min(GetSize(), InSize);
		m_DataEnd = zen::Min(m_DataEnd, m_Data + InSize);
	}

	/** Modifies the view to be the middle part by taking up to the given number of bytes from the given offset. */
	inline void MidInline(uint64_t InOffset, uint64_t InSize = ~uint64_t(0))
	{
		RightChopInline(InOffset);
		if (InSize != ~uint64_t(0))
		{
			LeftInline(InSize);
		}
	}

	/** Returns the middle part of the view by taking up to the given number of bytes from the given position. */
	[[nodiscard]] inline MemoryView Mid(uint64_t InOffset, uint64_t InSize = ~uint64_t(0)) const
	{
		MemoryView View(*this);
		View.MidInline(InOffset, InSize);
		return View;
	}

	constexpr void Reset()
	{
		m_Data	  = nullptr;
		m_DataEnd = nullptr;
	}

private:
	const uint8_t* m_Data	 = nullptr;
	const uint8_t* m_DataEnd = nullptr;

	/** Returns the data pointer advanced by an offset in bytes. */
	inline constexpr const uint8_t* GetDataAtOffsetNoCheck(uint64_t InOffset) const { return m_Data + InOffset; }
};

inline MutableMemoryView
MutableMemoryView::CopyFrom(MemoryView InView) const
{
	ZEN_ASSERT(InView.GetSize() <= GetSize());
	memcpy(m_Data, InView.GetData(), InView.GetSize());
	return RightChop(InView.GetSize());
}

/** Advances the start of the view by an offset, which is clamped to stay within the view. */
inline MemoryView
operator+(const MemoryView& View, uint64_t Offset)
{
	return MemoryView(View) += Offset;
}

/** Advances the start of the view by an offset, which is clamped to stay within the view. */
inline MemoryView
operator+(uint64_t Offset, const MemoryView& View)
{
	return MemoryView(View) += Offset;
}

/** Advances the start of the view by an offset, which is clamped to stay within the view. */
inline MutableMemoryView
operator+(const MutableMemoryView& View, uint64_t Offset)
{
	return MutableMemoryView(View) += Offset;
}

/** Advances the start of the view by an offset, which is clamped to stay within the view. */
inline MutableMemoryView
operator+(uint64_t Offset, const MutableMemoryView& View)
{
	return MutableMemoryView(View) += Offset;
}

/**
 * Make a non-owning view of the memory of the initializer list.
 *
 * This overload is only available when the element type does not need to be deduced.
 */
template<typename T>
[[nodiscard]] inline MemoryView
MakeMemoryView(std::initializer_list<typename std::type_identity<T>::type> List)
{
	return MemoryView(List.begin(), List.size() * sizeof(T));
}

/** Make a non-owning view of the memory of the contiguous container. */
template<ContiguousRange R>
[[nodiscard]] constexpr inline MemoryView
MakeMemoryView(const R& Container)
{
	std::span Span = Container;
	return MemoryView(Span.data(), Span.size_bytes());
}

/** Make a non-owning const view starting at Data and ending at DataEnd. */

[[nodiscard]] inline MemoryView
MakeMemoryView(const void* Data, const void* DataEnd)
{
	return MemoryView(Data, DataEnd);
}

[[nodiscard]] inline MemoryView
MakeMemoryView(const void* Data, uint64_t Size)
{
	return MemoryView(Data, reinterpret_cast<const uint8_t*>(Data) + Size);
}

/**
 * Make a non-owning mutable view of the memory of the initializer list.
 *
 * This overload is only available when the element type does not need to be deduced.
 */
template<typename T>
[[nodiscard]] inline MutableMemoryView
MakeMutableMemoryView(std::initializer_list<typename std::type_identity<T>::type> List)
{
	return MutableMemoryView(List.begin(), List.size() * sizeof(T));
}

/** Make a non-owning mutable view of the memory of the contiguous container. */
template<ContiguousRange R>
[[nodiscard]] constexpr inline MutableMemoryView
MakeMutableMemoryView(R& Container)
{
	std::span Span = Container;
	return MutableMemoryView(Span.data(), Span.size_bytes());
}

/** Make a non-owning mutable view starting at Data and ending at DataEnd. */

[[nodiscard]] inline MutableMemoryView
MakeMutableMemoryView(void* Data, void* DataEnd)
{
	return MutableMemoryView(Data, DataEnd);
}

void memory_forcelink();  // internal

}  // namespace zen
